{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "name": "VMAS: RLlib.ipynb",
   "provenance": [],
   "collapsed_sections": [
    "0NsC_EwfCF5I"
   ]
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU",
  "gpuClass": "standard"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "## Initialization"
   ],
   "metadata": {
    "id": "0NsC_EwfCF5I"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cP9ijqwvIXGd",
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "! git clone https://github.com/proroklab/VectorizedMultiAgentSimulator.git\n",
    "%cd /content/VectorizedMultiAgentSimulator\n",
    "!pip install -e .\n",
    "!pip install \"ray[rllib]\"==2.2 wandb\n",
    "!pip install \"pydantic<2\" numpy==1.23.5"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "#@title\n",
    "!sudo apt-get update\n",
    "!sudo apt-get install python3-opengl xvfb\n",
    "!pip install pyvirtualdisplay\n",
    "import pyvirtualdisplay\n",
    "display = pyvirtualdisplay.Display(visible=False, size=(1400, 900))\n",
    "display.start()"
   ],
   "metadata": {
    "id": "zjnXLxaOMLuv",
    "pycharm": {
     "is_executing": true
    }
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Run\n"
   ],
   "metadata": {
    "id": "jAAA3DXGCLkF"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import os\n",
    "has_gpu = torch.cuda.is_available()\n",
    "os.environ[\"RLLIB_NUM_GPUS\"] = \"1\" if has_gpu else \"0\""
   ],
   "metadata": {
    "id": "x6Kknt4W8gym"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "#  Copyright (c) 2022.\n",
    "#  ProrokLab (https://www.proroklab.org/)\n",
    "#  All rights reserved.\n",
    "\n",
    "import os\n",
    "from typing import Dict, Optional\n",
    "\n",
    "import numpy as np\n",
    "import ray\n",
    "from ray import tune\n",
    "from ray.rllib import BaseEnv, Policy, RolloutWorker\n",
    "from ray.rllib.algorithms.callbacks import DefaultCallbacks\n",
    "from ray.rllib.agents.ppo import PPOTrainer\n",
    "from ray.rllib.evaluation import Episode, MultiAgentEpisode\n",
    "from ray.rllib.utils.typing import PolicyID\n",
    "from ray.tune import register_env\n",
    "\n",
    "import wandb\n",
    "from vmas import make_env, Wrapper\n",
    "\n",
    "scenario_name = \"balance\"\n",
    "\n",
    "# Scenario specific variables.\n",
    "# When modifying this also modify env_config and env_creator\n",
    "n_agents = 4\n",
    "\n",
    "# Common variables\n",
    "continuous_actions = True\n",
    "max_steps = 200\n",
    "num_vectorized_envs = 96\n",
    "num_workers = 1\n",
    "vmas_device = \"cpu\"  # or cuda\n",
    "\n",
    "\n",
    "def env_creator(config: Dict):\n",
    "    env = make_env(\n",
    "        scenario=config[\"scenario_name\"],\n",
    "        num_envs=config[\"num_envs\"],\n",
    "        device=config[\"device\"],\n",
    "        continuous_actions=config[\"continuous_actions\"],\n",
    "        wrapper=Wrapper.RLLIB,\n",
    "        max_steps=config[\"max_steps\"],\n",
    "        # Scenario specific variables\n",
    "        n_agents=config[\"n_agents\"],\n",
    "    )\n",
    "    return env\n",
    "\n",
    "\n",
    "if not ray.is_initialized():\n",
    "    ray.init()\n",
    "    print(\"Ray init!\")\n",
    "register_env(scenario_name, lambda config: env_creator(config))\n",
    "\n",
    "\n",
    "class EvaluationCallbacks(DefaultCallbacks):\n",
    "    def on_episode_step(\n",
    "        self,\n",
    "        *,\n",
    "        worker: RolloutWorker,\n",
    "        base_env: BaseEnv,\n",
    "        episode: MultiAgentEpisode,\n",
    "        **kwargs,\n",
    "    ):\n",
    "        info = episode.last_info_for()\n",
    "        for a_key in info.keys():\n",
    "            for b_key in info[a_key]:\n",
    "                try:\n",
    "                    episode.user_data[f\"{a_key}/{b_key}\"].append(info[a_key][b_key])\n",
    "                except KeyError:\n",
    "                    episode.user_data[f\"{a_key}/{b_key}\"] = [info[a_key][b_key]]\n",
    "\n",
    "    def on_episode_end(\n",
    "        self,\n",
    "        *,\n",
    "        worker: RolloutWorker,\n",
    "        base_env: BaseEnv,\n",
    "        policies: Dict[str, Policy],\n",
    "        episode: MultiAgentEpisode,\n",
    "        **kwargs,\n",
    "    ):\n",
    "        info = episode.last_info_for()\n",
    "        for a_key in info.keys():\n",
    "            for b_key in info[a_key]:\n",
    "                metric = np.array(episode.user_data[f\"{a_key}/{b_key}\"])\n",
    "                episode.custom_metrics[f\"{a_key}/{b_key}\"] = np.sum(metric).item()\n",
    "\n",
    "\n",
    "class RenderingCallbacks(DefaultCallbacks):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.frames = []\n",
    "\n",
    "    def on_episode_step(\n",
    "        self,\n",
    "        *,\n",
    "        worker: RolloutWorker,\n",
    "        base_env: BaseEnv,\n",
    "        policies: Optional[Dict[PolicyID, Policy]] = None,\n",
    "        episode: Episode,\n",
    "        **kwargs,\n",
    "    ) -> None:\n",
    "        self.frames.append(base_env.vector_env.try_render_at(mode=\"rgb_array\"))\n",
    "\n",
    "    def on_episode_end(\n",
    "        self,\n",
    "        *,\n",
    "        worker: RolloutWorker,\n",
    "        base_env: BaseEnv,\n",
    "        policies: Dict[PolicyID, Policy],\n",
    "        episode: Episode,\n",
    "        **kwargs,\n",
    "    ) -> None:\n",
    "        vid = np.transpose(self.frames, (0, 3, 1, 2))\n",
    "        episode.media[\"rendering\"] = wandb.Video(vid, fps=1 / base_env.vector_env.env.world.dt, format=\"mp4\")\n",
    "        self.frames = []\n",
    "\n",
    "\n",
    "def train():\n",
    "\n",
    "    RLLIB_NUM_GPUS = int(os.environ.get(\"RLLIB_NUM_GPUS\", \"0\"))\n",
    "    num_gpus = 0.001 if RLLIB_NUM_GPUS > 0 else 0  # Driver GPU\n",
    "    num_gpus_per_worker = (\n",
    "        (RLLIB_NUM_GPUS - num_gpus) / (num_workers + 1) if vmas_device == \"cuda\" else 0\n",
    "    )\n",
    "\n",
    "    tune.run(\n",
    "        PPOTrainer,\n",
    "        stop={\"training_iteration\": 5000},\n",
    "        checkpoint_freq=1,\n",
    "        keep_checkpoints_num=2,\n",
    "        checkpoint_at_end=True,\n",
    "        checkpoint_score_attr=\"episode_reward_mean\",\n",
    "        # callbacks=[\n",
    "        #     WandbLoggerCallback(\n",
    "        #        project=f\"{scenario_name}\",\n",
    "        #        api_key=\"\",\n",
    "        #    )\n",
    "        # ],\n",
    "        config={\n",
    "            \"seed\": 0,\n",
    "            \"framework\": \"torch\",\n",
    "            \"env\": scenario_name,\n",
    "            \"kl_coeff\": 0.01,\n",
    "            \"kl_target\": 0.01,\n",
    "            \"lambda\": 0.9,\n",
    "            \"clip_param\": 0.2,\n",
    "            \"vf_loss_coeff\": 1,\n",
    "            \"vf_clip_param\": float(\"inf\"),\n",
    "            \"entropy_coeff\": 0,\n",
    "            \"train_batch_size\": 60000,\n",
    "            \"rollout_fragment_length\": 125,\n",
    "            \"sgd_minibatch_size\": 4096,\n",
    "            \"num_sgd_iter\": 40,\n",
    "            \"num_gpus\": num_gpus,\n",
    "            \"num_workers\": num_workers,\n",
    "            \"num_gpus_per_worker\": num_gpus_per_worker,\n",
    "            \"num_envs_per_worker\": num_vectorized_envs,\n",
    "            \"lr\": 5e-5,\n",
    "            \"gamma\": 0.99,\n",
    "            \"use_gae\": True,\n",
    "            \"use_critic\": True,\n",
    "            \"batch_mode\": \"truncate_episodes\",\n",
    "            \"env_config\": {\n",
    "                \"device\": vmas_device,\n",
    "                \"num_envs\": num_vectorized_envs,\n",
    "                \"scenario_name\": scenario_name,\n",
    "                \"continuous_actions\": continuous_actions,\n",
    "                \"max_steps\": max_steps,\n",
    "                # Scenario specific variables\n",
    "                \"n_agents\": n_agents,\n",
    "            },\n",
    "            \"evaluation_interval\": 5,\n",
    "            \"evaluation_duration\": 1,\n",
    "            \"evaluation_num_workers\": 0,\n",
    "            \"evaluation_parallel_to_training\": False,\n",
    "            \"evaluation_config\": {\n",
    "                \"num_envs_per_worker\": 1,\n",
    "                \"env_config\": {\n",
    "                    \"num_envs\": 1,\n",
    "                },\n",
    "                # \"callbacks\": MultiCallbacks([RenderingCallbacks, EvaluationCallbacks]),\n",
    "            },\n",
    "            \"callbacks\": EvaluationCallbacks,\n",
    "        },\n",
    "    )\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    train()"
   ],
   "metadata": {
    "id": "DFsqhhL97JFh"
   },
   "execution_count": null,
   "outputs": []
  }
 ]
}
